package org.zjx.matcher;

import lombok.Data;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zjx.db.ItemTemplateDB;
import org.zjx.model.prop.ItemTemplate;
import org.zjx.util.ImageUtil;

import java.awt.image.BufferedImage;
import java.util.*;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;

public class TemplateMatcher {
    private static final Logger logger = LoggerFactory.getLogger(TemplateMatcher.class);

    private final ItemTemplateDB templateDB;
    private final List<TemplateData> templateDataList;
    private final ExecutorService executorService;

    public TemplateMatcher(ItemTemplateDB templateDB) {
        this.templateDB = Objects.requireNonNull(templateDB);
        this.executorService = Executors.newWorkStealingPool();
        this.templateDataList = new CopyOnWriteArrayList<>();
        preloadTemplates();
    }

    private void preloadTemplates() {
        templateDB.findAll().parallelStream().forEach(template -> {
            try {
                if (template.getImageBase64() != null && !template.getImageBase64().isEmpty()) {
                    BufferedImage img = ImageUtil.stringToBufferedImage(template.getImageBase64());
                    // 预计算模板的非透明像素信息
                    TemplatePixelData pixelData = extractTemplatePixelData(img);
                    templateDataList.add(new TemplateData(template, img, pixelData));
                }
            } catch (Exception e) {
                logger.error("模板预加载失败 ID={}: {}", template.getId(), e.getMessage());
            }
        });
        logger.info("已预加载 {} 个模板", templateDataList.size());
    }

    public Optional<ItemTemplate> match(BufferedImage input) {
        long startTime = System.currentTimeMillis();
        try {
            Optional<TemplateData> matched = templateDataList.parallelStream()
                    .filter(templateData -> isMatch(input, templateData))
                    .findFirst();

            logger.debug("匹配耗时: {}ms", System.currentTimeMillis() - startTime);
            return matched.map(td -> templateDB.findById(td.template.getId()));
        } catch (Exception e) {
            logger.error("图像匹配失败: {}", e.getMessage());
            return Optional.empty();
        }
    }

    private boolean isMatch(BufferedImage input, TemplateData templateData) {
        TemplatePixelData pixelData = templateData.pixelData;
        int templateWidth = templateData.image.getWidth();
        int templateHeight = templateData.image.getHeight();

        // 输入图像比模板小，直接不匹配
        if (input.getWidth() < templateWidth || input.getHeight() < templateHeight) {
            return false;
        }

        // 使用滑动窗口检查所有可能位置
        for (int y = 0; y <= input.getHeight() - templateHeight; y++) {
            for (int x = 0; x <= input.getWidth() - templateWidth; x++) {
                if (checkPosition(input, x, y, pixelData)) {
                    return true;
                }
            }
        }
        return false;
    }

    // 检查特定位置是否匹配
    private boolean checkPosition(BufferedImage input, int startX, int startY, TemplatePixelData pixelData) {
        for (PixelInfo pixel : pixelData.pixels) {
            int inputX = startX + pixel.x;
            int inputY = startY + pixel.y;
            if (input.getRGB(inputX, inputY) != pixel.rgb) {
                return false; // 快速失败
            }
        }
        return true;
    }

    // 预计算模板的非透明像素信息
    private TemplatePixelData extractTemplatePixelData(BufferedImage template) {
        List<PixelInfo> pixels = new ArrayList<>();
        for (int y = 0; y < template.getHeight(); y++) {
            for (int x = 0; x < template.getWidth(); x++) {
                int rgb = template.getRGB(x, y);
                if ((rgb >>> 24) != 0x00) { // 非透明像素
                    pixels.add(new PixelInfo(x, y, rgb));
                }
            }
        }
        return new TemplatePixelData(pixels);
    }

    // 批量匹配优化
    public List<ItemTemplate> matchBatch(List<BufferedImage> inputs) {
        return inputs.parallelStream()
                .map(input -> {
                    Optional<ItemTemplate> result = match(input);
                    return result.orElse(null); // 未匹配的返回null
                })
                .collect(Collectors.toList());
    }

    public void shutdown() {
        executorService.shutdown();
        templateDataList.clear();
    }

    // 模板数据类
    private static class TemplateData {
        final ItemTemplate template;
        final BufferedImage image;
        final TemplatePixelData pixelData;

        TemplateData(ItemTemplate template, BufferedImage image, TemplatePixelData pixelData) {
            this.template = template;
            this.image = image;
            this.pixelData = pixelData;
        }
    }

    // 模板像素数据
    @Data
    private static class TemplatePixelData {
        private List<PixelInfo> pixels;

        public TemplatePixelData(List<PixelInfo> pixels) {
            this.pixels = Collections.unmodifiableList(pixels);
        }
    }

    @Data
    // 像素信息
    private static class PixelInfo {
        private final int x;
        private final int y;
        private final int rgb;

        public PixelInfo(int x, int y, int rgb) {
            this.x = x;
            this.y = y;
            this.rgb = rgb;
        }
    }
}